#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
营业执照OCR性能配置模块
提供针对营业执照识别优化的OCR参数配置
"""

import logging
from typing import Dict, Any, Optional
from dataclasses import dataclass


@dataclass
class OCRPerformanceConfig:
    """OCR性能配置"""
    # 检测参数
    det_db_thresh: float = 0.3          # 文本检测阈值
    det_db_box_thresh: float = 0.6      # 文本框阈值
    det_db_unclip_ratio: float = 1.5    # 文本框扩展比例
    det_max_side_len: int = 960         # 检测图像最大边长
    
    # 识别参数
    rec_batch_num: int = 6              # 识别批处理大小
    max_text_length: int = 25           # 最大文本长度
    
    # 方向分类参数
    use_angle_cls: bool = True          # 是否使用方向分类
    cls_thresh: float = 0.9             # 方向分类阈值
    
    # 性能参数
    use_gpu: bool = False               # 是否使用GPU
    enable_mkldnn: bool = True          # 是否启用MKLDNN加速
    cpu_threads: int = 4                # CPU线程数
    
    # 内存优化参数
    use_tensorrt: bool = False          # 是否使用TensorRT
    precision: str = "fp32"             # 精度模式: fp32, fp16, int8
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'det_db_thresh': self.det_db_thresh,
            'det_db_box_thresh': self.det_db_box_thresh,
            'det_db_unclip_ratio': self.det_db_unclip_ratio,
            'det_max_side_len': self.det_max_side_len,
            'rec_batch_num': self.rec_batch_num,
            'max_text_length': self.max_text_length,
            'use_angle_cls': self.use_angle_cls,
            'cls_thresh': self.cls_thresh,
            'use_gpu': self.use_gpu,
            'enable_mkldnn': self.enable_mkldnn,
            'cpu_threads': self.cpu_threads,
            'use_tensorrt': self.use_tensorrt,
            'precision': self.precision
        }


class OCRConfigManager:
    """OCR配置管理器"""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
        
        # 预定义的配置模式
        self._config_modes = {
            'fast': self._get_fast_config(),
            'balanced': self._get_balanced_config(),
            'accurate': self._get_accurate_config(),
            'memory_optimized': self._get_memory_optimized_config()
        }
        
        self._current_mode = 'balanced'
        self._current_config = self._config_modes[self._current_mode]
    
    def _get_fast_config(self) -> OCRPerformanceConfig:
        """快速模式配置 - 优先速度"""
        return OCRPerformanceConfig(
            det_db_thresh=0.4,
            det_db_box_thresh=0.7,
            det_db_unclip_ratio=1.2,
            det_max_side_len=640,
            rec_batch_num=8,
            max_text_length=20,
            use_angle_cls=False,  # 关闭方向分类以提高速度
            cls_thresh=0.9,
            use_gpu=False,
            enable_mkldnn=True,
            cpu_threads=2,
            use_tensorrt=False,
            precision="fp32"
        )
    
    def _get_balanced_config(self) -> OCRPerformanceConfig:
        """平衡模式配置 - 速度和精度平衡"""
        return OCRPerformanceConfig(
            det_db_thresh=0.3,
            det_db_box_thresh=0.6,
            det_db_unclip_ratio=1.5,
            det_max_side_len=960,
            rec_batch_num=6,
            max_text_length=25,
            use_angle_cls=True,
            cls_thresh=0.9,
            use_gpu=False,
            enable_mkldnn=True,
            cpu_threads=4,
            use_tensorrt=False,
            precision="fp32"
        )
    
    def _get_accurate_config(self) -> OCRPerformanceConfig:
        """精确模式配置 - 优先精度"""
        return OCRPerformanceConfig(
            det_db_thresh=0.2,
            det_db_box_thresh=0.5,
            det_db_unclip_ratio=2.0,
            det_max_side_len=1280,
            rec_batch_num=4,
            max_text_length=30,
            use_angle_cls=True,
            cls_thresh=0.8,
            use_gpu=False,
            enable_mkldnn=True,
            cpu_threads=6,
            use_tensorrt=False,
            precision="fp32"
        )
    
    def _get_memory_optimized_config(self) -> OCRPerformanceConfig:
        """内存优化配置 - 减少内存使用"""
        return OCRPerformanceConfig(
            det_db_thresh=0.35,
            det_db_box_thresh=0.65,
            det_db_unclip_ratio=1.3,
            det_max_side_len=480,  # 较小的图像尺寸
            rec_batch_num=2,       # 较小的批处理大小
            max_text_length=20,
            use_angle_cls=False,   # 关闭方向分类节省内存
            cls_thresh=0.9,
            use_gpu=False,
            enable_mkldnn=False,   # 关闭MKLDNN节省内存
            cpu_threads=2,         # 较少的线程数
            use_tensorrt=False,
            precision="fp32"
        )
    
    def set_mode(self, mode: str) -> bool:
        """设置配置模式"""
        if mode not in self._config_modes:
            self.logger.warning(f"未知的配置模式: {mode}, 可用模式: {list(self._config_modes.keys())}")
            return False
        
        self._current_mode = mode
        self._current_config = self._config_modes[mode]
        self.logger.info(f"OCR配置模式已切换到: {mode}")
        return True
    
    def get_current_config(self) -> OCRPerformanceConfig:
        """获取当前配置"""
        return self._current_config
    
    def get_current_mode(self) -> str:
        """获取当前模式"""
        return self._current_mode
    
    def get_available_modes(self) -> list:
        """获取可用的配置模式"""
        return list(self._config_modes.keys())
    
    def update_config(self, **kwargs) -> bool:
        """更新当前配置的特定参数"""
        try:
            config_dict = self._current_config.to_dict()
            
            # 更新指定的参数
            for key, value in kwargs.items():
                if key in config_dict:
                    setattr(self._current_config, key, value)
                    self.logger.debug(f"更新OCR配置参数: {key} = {value}")
                else:
                    self.logger.warning(f"未知的配置参数: {key}")
                    return False
            
            return True
            
        except Exception as e:
            self.logger.error(f"更新OCR配置失败: {str(e)}")
            return False
    
    def get_paddle_ocr_params(self) -> Dict[str, Any]:
        """获取PaddleOCR初始化参数"""
        config = self._current_config
        return {
            'use_angle_cls': config.use_angle_cls,
            'use_gpu': config.use_gpu,
            'lang': 'ch',
            'show_log': False,
            'det_db_thresh': config.det_db_thresh,
            'det_db_box_thresh': config.det_db_box_thresh,
            'det_db_unclip_ratio': config.det_db_unclip_ratio,
            'det_max_side_len': config.det_max_side_len,
            'enable_mkldnn': config.enable_mkldnn,
            'cpu_threads': config.cpu_threads,
            'rec_batch_num': config.rec_batch_num,
            'max_text_length': config.max_text_length,
            'cls_thresh': config.cls_thresh,
            'use_tensorrt': config.use_tensorrt,
            'precision': config.precision
        }
    
    def get_rapid_ocr_params(self) -> Dict[str, Any]:
        """获取RapidOCR初始化参数"""
        config = self._current_config
        return {
            'text_score': config.det_db_thresh,
            'use_angle_cls': config.use_angle_cls,
            'use_text_det': True,
            'use_text_rec': True,
            'print_verbose': False,
            'min_height': 30,
            'width_height_ratio': 8
        }
    
    def auto_adjust_for_image_size(self, image_width: int, image_height: int):
        """根据图像大小自动调整配置"""
        max_dimension = max(image_width, image_height)
        
        if max_dimension > 2000:
            # 大图像，使用内存优化配置
            self.set_mode('memory_optimized')
            self.logger.info(f"检测到大图像 ({max_dimension}px)，切换到内存优化模式")
        elif max_dimension < 800:
            # 小图像，使用精确模式
            self.set_mode('accurate')
            self.logger.info(f"检测到小图像 ({max_dimension}px)，切换到精确模式")
        else:
            # 中等大小图像，使用平衡模式
            self.set_mode('balanced')
            self.logger.info(f"检测到中等图像 ({max_dimension}px)，使用平衡模式")
    
    def auto_adjust_for_performance(self, avg_processing_time: float, memory_usage_mb: float):
        """根据性能指标自动调整配置"""
        # 如果处理时间过长（超过5秒），切换到快速模式
        if avg_processing_time > 5.0:
            if self._current_mode != 'fast':
                self.set_mode('fast')
                self.logger.info(f"检测到处理时间过长 ({avg_processing_time:.2f}s)，切换到快速模式")
        
        # 如果内存使用过高（超过800MB），切换到内存优化模式
        elif memory_usage_mb > 800:
            if self._current_mode != 'memory_optimized':
                self.set_mode('memory_optimized')
                self.logger.info(f"检测到内存使用过高 ({memory_usage_mb:.1f}MB)，切换到内存优化模式")
        
        # 如果性能良好，使用平衡模式
        elif avg_processing_time < 3.0 and memory_usage_mb < 400:
            if self._current_mode not in ['balanced', 'accurate']:
                self.set_mode('balanced')
                self.logger.info("性能良好，切换到平衡模式")


# 全局OCR配置管理器实例
ocr_config_manager = OCRConfigManager()


def get_optimized_ocr_config(mode: str = 'balanced') -> Dict[str, Any]:
    """获取优化的OCR配置"""
    ocr_config_manager.set_mode(mode)
    return ocr_config_manager.get_paddle_ocr_params()


def auto_tune_ocr_config(image_width: int, image_height: int, 
                        avg_processing_time: float = 0.0, 
                        memory_usage_mb: float = 0.0) -> Dict[str, Any]:
    """自动调优OCR配置"""
    # 根据图像大小调整
    ocr_config_manager.auto_adjust_for_image_size(image_width, image_height)
    
    # 根据性能指标调整
    if avg_processing_time > 0 or memory_usage_mb > 0:
        ocr_config_manager.auto_adjust_for_performance(avg_processing_time, memory_usage_mb)
    
    return ocr_config_manager.get_paddle_ocr_params()